from kalman_filter import kalman_filter_2D
import matplotlib.pyplot as plt
import numpy as np

kf2 = kalman_filter_2D(1, 0.01, 25, 25)

## 数据生成

## 1. 正余弦组合
# z_x = 5 * np.sin(np.arange(0, 10, 0.1))
# z_y = np.arange(0, 10, 0.1)
# z_x_noise = z_x + np.random.normal(0, 1, size=(100,))
# z_y_noise = z_y + np.random.normal(0, 1, size=(100,))
# Z_x = iter(z_x_noise)
# Z_y = iter(z_y_noise)

# 2. 均匀加速的直线运动
z_x = np.arange(0, 100, 1)
z_y = 0.01 * np.arange(0, 100, 1) ** 2
z_x_noise = z_x + 5 * np.random.normal(0, 1, size=(100,))
z_y_noise = z_y + 5 * np.random.normal(0, 1, size=(100,))
Z_x = iter(z_x_noise)
Z_y = iter(z_y_noise)

## end

# 卡尔曼滤波
pred_x, pred_y = kf2.iterate(100, Z_x, Z_y)

# 绘图
plt.plot(z_x, z_y, label="True")
plt.plot(z_x_noise, z_y_noise, "r", label="Measure")
plt.plot(pred_x, pred_y, "g", label="Predict")
plt.legend()
plt.axis("off")
plt.show()
